import numpy as np
import os
import itertools
from tqdm import trange


class unstructuredGridContainer(object):

    def __init__(self, **kwargs):
        self.origin = kwargs.get('origin', [0, 0, 0])
        self.bound = kwargs.get('bound', [127, 127, 127])
        self.resolution = kwargs.get('resolution', [1, 1, 1])
        self.type = kwargs.get('type', 12)
        self.wd = kwargs.get('wd', os.getcwd())
        self.filename = kwargs.get('filename', '')

        self.numNodes = 0
        self.numElements = 0
        self.f = None

    def __str__(self):
        return 'Unstructured Grid Container(origin={}, bound={}, resolution={}, type={:d})'.format(self.origin,
                                                                                                   self.bound,
                                                                                                   self.resolution,
                                                                                                   self.type)

    def __repr__(self):
        return 'unstructuredGridContainer Object'

    def genMesh(self, **kwargs):

        tocsv = kwargs.get('tocsv', False)

        xNumNodes = (self.bound[0] - self.origin[0] + 1) // self.resolution[0]
        yNumNodes = (self.bound[1] - self.origin[1] + 1) // self.resolution[1]
        zNumNodes = (self.bound[2] - self.origin[2] + 1) // self.resolution[2]

        xCoord = np.linspace(self.origin[0], self.bound[0], xNumNodes)
        yCoord = np.linspace(self.origin[1], self.bound[1], yNumNodes)
        zCoord = np.linspace(self.origin[2], self.bound[2], zNumNodes)

        numNodes = xNumNodes * yNumNodes * zNumNodes
        numElements = (xNumNodes - 1) * (yNumNodes - 1) * (zNumNodes - 1)

        print('正在创建网格......')
        print('共有节点 {:,} 个，共有单元 {:,} 个'.format(numNodes, numElements))

        nodeArray = []
        for combination in itertools.product(xCoord, yCoord, zCoord):
            nodeArray.append(combination)
        nodeArray = np.array(nodeArray)
        nodeArray[:] = nodeArray[:, [2, 1, 0]]

        if self.type == 12:
            elementArray = np.zeros([numElements, 8])
        _arr = unstructuredGridContainer.connectivityGenerator(xNumNodes, yNumNodes, zNumNodes)
        for i, _ in enumerate(_arr):
            elementArray[i] = _
        elementArray = elementArray.astype(int)

        print('网格创建成功')

        if tocsv:
            np.savetxt(os.path.join(self.wd, '{}_node.csv'.format(self.filename)), nodeArray, delimiter='\t',
                       fmt='%1.3f')
            np.savetxt(os.path.join(self.wd, '{}_element.csv'.format(self.filename)), elementArray, delimiter='\t',
                       fmt='%d')

        return nodeArray, elementArray

    def loadMesh(self, **kwargs):

        meshSource = kwargs.get('meshSource', 'inp')

        pass

    def fetchData(self):
        pass

    def writeMesh(self, **kwargs):
        self.f = open(os.path.join(self.wd, '{}.vtk'.format(self.filename)), 'w', newline='\n')
        self.f.write('# vtk DataFile Version 3.0\n')
        self.f.write('Write by mmtoolkit\n')
        self.f.write('ASCII\n')
        self.f.write('DATASET UNSTRUCTURED_GRID\n')

        nodeArray = kwargs.get('node', np.array(()))
        elementArray = kwargs.get('element', np.array(()))

        self.numNodes = nodeArray.shape[0]
        self.numElements = elementArray.shape[0]

        tempArray = np.zeros([self.numElements, 1])
        tempArray[:] = elementArray.shape[1]
        elementArray = np.hstack((tempArray, elementArray))

        tempArray[:] = self.type

        self.f.write('POINTS {} double\n'.format(self.numNodes))
        np.savetxt(self.f, nodeArray, fmt='%.5f')
        self.f.write('CELLS {} {}\n'.format(self.numElements, elementArray.size))
        np.savetxt(self.f, elementArray, fmt='%d')
        self.f.write('CELL_TYPES {}\n'.format(self.numElements))
        np.savetxt(self.f, tempArray, fmt='%d')

    def writePointData(self, **kwargs):

        tensorVar = kwargs.get('tensor', None)
        vectorVar = kwargs.get('vector', None)
        scalarVar = kwargs.get('scalar', None)
        prefix = kwargs.get('prefix', None)
        suffix = kwargs.get('suffix', None)
        fmt = kwargs.get('fmt', 'csv')

        self.f.write('POINT_DATA {}\n'.format(self.numNodes))

        if len(tensorVar) == 0:
            pass
        else:
            for tVar in tensorVar:
                tensorArray = np.loadtxt(os.path.join(self.wd, '{}_{}_{}.{}'.format(prefix, tVar, suffix, fmt)),
                                         delimiter='\t')
                tensorArray = unstructuredGridContainer.tensorExpand(tensorArray, self.numElements)
                self.f.write('TENSORS {} double\n'.format(tVar))
                for ndArray in tensorArray:
                    np.savetxt(self.f, ndArray, fmt='%.5f')

        if len(vectorVar) == 0:
            pass
        else:
            for vVar in vectorVar:
                if vVar == 'ORI':
                    vectorArray = np.loadtxt(os.path.join(self.wd, '{}_{}_{}.{}'.format(prefix, vVar, suffix, fmt)),
                                             delimiter='\t')
                    vectorArray = vectorArray[:, -3:]
                else:
                    vectorArray = np.loadtxt(os.path.join(self.wd, '{}_{}_{}.{}'.format(prefix, vVar, suffix, fmt)),
                                             delimiter='\t')

                self.f.write('VECTORS {} double\n'.format(vVar))
                np.savetxt(self.f, vectorArray, fmt='%.5f')

        if len(scalarVar) == 0:
            pass
        else:
            for sVar in scalarVar:
                scalarArray = np.loadtxt(os.path.join(self.wd, '{}_{}_{}.{}'.format(prefix, sVar, suffix, fmt)),
                                         delimiter='\t')

                self.f.write('SCALARS {} double 1\n'.format(sVar))
                self.f.write('LOOKUP_TABLE default\n')
                np.savetxt(self.f, scalarArray, fmt='%.5f')

    def writeCellData(self, **kwargs):

        tensorVar = kwargs.get('tensor', None)
        vectorVar = kwargs.get('vector', None)
        scalarVar = kwargs.get('scalar', None)
        prefix = kwargs.get('prefix', None)
        suffix = kwargs.get('suffix', None)
        fmt = kwargs.get('fmt', 'csv')

        self.f.write('CELL_DATA {}\n'.format(self.numElements))

        if len(tensorVar) == 0:
            pass
        else:
            for tVar in tensorVar:
                tensorArray = np.loadtxt(os.path.join(self.wd, '{}_{}_{}.{}'.format(prefix, tVar, suffix, fmt)),
                                         delimiter='\t')
                tensorArray = unstructuredGridContainer.tensorExpand(tensorArray, self.numElements)
                self.f.write('TENSORS {} double\n'.format(tVar))
                for ndArray in tensorArray:
                    np.savetxt(self.f, ndArray, fmt='%.5f')

        if len(vectorVar) == 0:
            pass
        else:
            for vVar in vectorVar:
                if vVar == 'ORI':
                    vectorArray = np.loadtxt(os.path.join(self.wd, '{}_{}_{}.{}'.format(prefix, vVar, suffix, fmt)),
                                             delimiter='\t')
                    vectorArray = vectorArray[:, -3:]
                else:
                    vectorArray = np.loadtxt(os.path.join(self.wd, '{}_{}_{}.{}'.format(prefix, vVar, suffix, fmt)),
                                             delimiter='\t')

                self.f.write('VECTORS {} double\n'.format(vVar))
                np.savetxt(self.f, vectorArray, fmt='%.5f')

        if len(scalarVar) == 0:
            pass
        else:
            for sVar in scalarVar:
                scalarArray = np.loadtxt(os.path.join(self.wd, '{}_{}_{}.{}'.format(prefix, sVar, suffix, fmt)),
                                         delimiter='\t')

                self.f.write('SCALARS {} double 1\n'.format(sVar))
                self.f.write('LOOKUP_TABLE default\n')
                np.savetxt(self.f, scalarArray, fmt='%.5f')

    @property
    def close(self):
        self.f.close()

    @staticmethod
    def connectivityGenerator(_x, _y, _z):
        for k in trange(_z - 1):
            for j in range(_y - 1):
                for i in range(_x - 1):
                    _con = np.array(
                        [k * _x * _y + j * _x + i,
                         k * _x * _y + j * _x + i + 1,
                         k * _x * _y + j * _x + _x + i + 1,
                         k * _x * _y + j * _x + _x + i,
                         k * _x * _y + j * _x + _x * _y + i,
                         k * _x * _y + j * _x + _x * _y + i + 1,
                         k * _x * _y + j * _x + _x * _y + _x + i + 1,
                         k * _x * _y + j * _x + _x * _y + _x + i,
                         ])
                    yield _con

    @staticmethod
    def tensorExpand(inTensorArray, n):
        outTensorArray = np.zeros([n, 9])
        outTensorArray[:, 0] = inTensorArray[:, 0]
        outTensorArray[:, 1] = inTensorArray[:, 3]
        outTensorArray[:, 2] = inTensorArray[:, 4]
        outTensorArray[:, 3] = inTensorArray[:, 3]
        outTensorArray[:, 4] = inTensorArray[:, 1]
        outTensorArray[:, 5] = inTensorArray[:, 5]
        outTensorArray[:, 6] = inTensorArray[:, 4]
        outTensorArray[:, 7] = inTensorArray[:, 5]
        outTensorArray[:, 8] = inTensorArray[:, 2]

        outTensorArray = outTensorArray.reshape([n, 3, 3])

        return outTensorArray
